package cn.codenest.listener;

import cn.codenest.db.DemoDB;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

/**
 * @author ：Hyman
 * @date ：Created in 2021/7/12 18:52
 * @description：可以自定义偏移量的再均衡监听器
 * @modified By：
 * @version: $
 */
public class SaveOffsetsOnReblance implements ConsumerRebalanceListener {

    public SaveOffsetsOnReblance(Map<TopicPartition, OffsetAndMetadata> currentOffsets, Consumer consumer) {
        this.currentOffsets = currentOffsets;
        this.consumer = consumer;
    }

    private Map<TopicPartition, OffsetAndMetadata> currentOffsets = new HashMap<>(128);
    private Consumer consumer;

    /**
     * create by: Hyman
     * description: 在再均衡开始之前，消费者停止读取消息之后被调用,如果在这里提交
     * 偏移量，那么下一个接管分区的消费者就知道从哪里开始读了
     * create time: 2021/7/9
     */
    @Override
    public void onPartitionsRevoked(Collection<TopicPartition> collection) {
        commitDBTransaction();
    }

    /**
     * create by: Hyman
     * description: 会在重新分配分区之后和消费者开始读取消息之前被调用
     * create time: 2021/7/9
     */
    @Override
    public void onPartitionsAssigned(Collection<TopicPartition> collection) {
        for (TopicPartition partition : collection) {
            consumer.seek(partition, getOffsetFromDB(partition));
        }
    }

    /**
     * create by: Hyman
     * description: 模拟将记录和偏移量放到同一个数据库事务里面放入数据库
     * create time: 2021/7/12
     */
    public void commitDBTransaction() {
        DemoDB.commitDBTransaction();
    }

    public long getOffsetFromDB(TopicPartition partition) {
        return DemoDB.offsets.get(
                DemoDB.getKey(partition.topic(),
                        partition.partition()));
    }


}
